package org.centsuse.intention_spring.invention.factory;

import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.centsuse.intention_spring.invention.config.models.ModelConfig;
import org.centsuse.intention_spring.invention.config.models.ModelConfiguration;
import org.centsuse.intention_spring.invention.utils.StringUtils;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class ModelFactory implements InitializingBean {

    /** 模型容器 */
    private final Map<String, ChatLanguageModel> modelRegistry = new ConcurrentHashMap<>();

    /** 通用对话模型 */
    private static final String DEFAULT_CHAT_MODEL = "ollama";

    /** 模型配置 */
    private final ModelConfiguration modelConfiguration;

    public ModelFactory(ModelConfiguration modelConfiguration) {
        this.modelConfiguration = modelConfiguration;
    }

    /**
     * @description 根据模型名称获取模型实体
     * @param modelKey 模型名称 yml中的ollama
     * @return ChatLanguageModel 模型实体
     */
    public ChatLanguageModel getModel(String modelKey){
        ChatLanguageModel model = modelRegistry.get(modelKey);
        if (model == null) {
            log.warn("[ModelFactory][getModel] Model not found: {}", modelKey);
        }
        return model;
    }

    /**
     * @description 根据模型名称获取模型实体
     * @return ChatLanguageModel 模型实体
     */
    public ChatLanguageModel getDefaultModel(){
        ChatLanguageModel model = modelRegistry.get(DEFAULT_CHAT_MODEL);
        if (model == null) {
            log.warn("[ModelFactory][getDefaultModel] Model not found: {}", DEFAULT_CHAT_MODEL);
        }
        return model;
    }


    @Override
    public void afterPropertiesSet() throws Exception {
        log.info("[ModelFactory][afterPropertiesSet] start initialization models, modelConfiguration {}", modelConfiguration);
        Map<String, ModelConfig> modelConfigs = modelConfiguration.getModelConfigMap();
        for (Map.Entry<String, ModelConfig> entry : modelConfigs.entrySet()) {
            if (StringUtils.isEmpty(entry.getKey()) || entry.getValue() == null) {
                continue;
            }
            if (modelRegistry.containsKey(entry.getKey())) {
                log.warn("[ModelFactory][getModel] Model {} already exists", entry.getKey());
            }
            // TODO 后续其他模型
            if (entry.getKey().contains("ollama")) {
                ChatLanguageModel chatLanguageModel = OllamaChatModel.builder()
                        .baseUrl(entry.getValue().getBaseUrl())
                        .modelName(entry.getValue().getModelName())
                        .temperature(entry.getValue().getTemperature())
                        .build();

                modelRegistry.put(entry.getKey(), chatLanguageModel);
            }
        }
        log.info("[ModelFactory][afterPropertiesSet] end initialization models, modelRegistry {}", modelRegistry);
    }
}
